/*
Copyright (c) 2008,2009, David Beck

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:

1. Redistributions of source code must retain the above copyright
   notice, this list of conditions and the following disclaimer.
2. Redistributions in binary form must reproduce the above copyright
   notice, this list of conditions and the following disclaimer in the
   documentation and/or other materials provided with the distribution.

THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*/

/**
   @file t__sec_simple.cc
   @brief implements a hand crafted security model to see how Rx works (used by the t1 tests)
   @see t__t1_client.cc
   @see t__t1_server.cc
*/

#define RT0D4_WINDOWS_H_INCLUDED
#include "t__sec_common.h"
#include "secobj_priv.hh"
#include "conn_priv.hh"
#include "crypt_buf.hh"
#include <string.h>
#include <assert.h>

using rt0d4::sec::crypt_buf;
using rt0d4::rx::conn_priv;
using rt0d4::rx::secobj_priv;

// rx/rx.h
// rx_connection_rx_lock

#define PRSTAT(CP,VAL) \
  if( CP->VAL() ) printf("  .. STAT %s : %ld\n", #VAL, CP->VAL() )

void print_rx_debug( struct rx_packet * apacket, struct rx_connection * aconn, char * prefix = "", struct rx_call * acall = NULL )
{
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  size_t xl = (apacket == NULL ? 0 : rx_GetDataSize(apacket));
  size_t hl = rx_GetSecurityHeaderSize(aconn);
  size_t ml = rx_GetSecurityMaxTrailerSize(aconn);

  printf(  "------ %s ------\n",prefix );
  printf(  "Len  : %d\n", xl );
  printf(  "Head : %d\n", hl );
  printf(  "Mac  : %d\n", ml );

  if( acall )
  {
    printf(  "Chan : %u\n", acall->channel );
    printf(  "nLeft: %u\n", acall->nLeft );
    printf(  "nFree: %u\n", acall->nFree );

    printf(  "Sent : %u\n", acall->bytesSent.low );
    printf(  "Recvd: %u\n", acall->bytesRcvd.low );
  }

  printf(  "#iovs: %d\n", (apacket == NULL ? 0 : apacket->niovecs) );
  for( size_t i=0;(apacket != NULL && i<1);++i )
  {
    printf("-- Internal --\n");
    printf("sz   : %d [%d]\n",apacket->wirevec[i].iov_len,i );
    printf("ptr  : %p [%d]\n",apacket->wirevec[i].iov_base,i );

    size_t ii=0;
    for( ;(ii<100 && ii<apacket->wirevec[i].iov_len);++ii )
      printf("%.2x", ((unsigned char *)apacket->wirevec[i].iov_base)[ii] );

    printf("\n");
  }

  for( size_t i=1;(apacket != NULL && i<apacket->niovecs);++i )
  {
    printf("-- User data --\n");
    printf("sz   : %d [%d]\n",apacket->wirevec[i].iov_len,i );
    printf("ptr  : %p [%d]\n",apacket->wirevec[i].iov_base,i );

    size_t ii=0;
    for( ;(ii<hl && ii<100 && ii<apacket->wirevec[i].iov_len);++ii )
      printf("%.2x", ((unsigned char *)apacket->wirevec[i].iov_base)[ii] );
    printf(" -> ");
    for( ;(ii<hl+xl && ii<100 && ii<apacket->wirevec[i].iov_len );++ii )
      printf("%.2x", ((unsigned char *)apacket->wirevec[i].iov_base)[ii] );
    printf(" -> ");
    for( ;(ii<hl+xl+ml && ii<100 && ii<apacket->wirevec[i].iov_len );++ii )
      printf("%.2x", ((unsigned char *)apacket->wirevec[i].iov_base)[ii] );
    printf(" -> ");
    for( ;(ii<apacket->wirevec[i].iov_len && ii<100 );++ii )
      printf("%.2x", ((unsigned char *)apacket->wirevec[i].iov_base)[ii] );
    printf( "\n" );
    printf(" <= ");
    for( ii=apacket->wirevec[i].iov_len-hl;ii<apacket->wirevec[i].iov_len;++ii )
      printf("%.2x", ((unsigned char *)apacket->wirevec[i].iov_base)[ii] );
    printf( "\n" );
  }
  printf(  "Sess : %s\n", (unsigned char *)cpriv->get_session_key() );

  PRSTAT(cpriv, n_prepare_packet_ok );
  PRSTAT(cpriv, n_send_packet_ok );
  PRSTAT(cpriv, n_check_auth_ok );
  PRSTAT(cpriv, n_create_challenge_ok );
  PRSTAT(cpriv, n_get_challenge_ok );
  PRSTAT(cpriv, n_get_response_ok );
  PRSTAT(cpriv, n_check_response_ok );
  PRSTAT(cpriv, n_check_packet_ok );
      /* */
  PRSTAT(cpriv, n_prepare_packet_err );
  PRSTAT(cpriv, n_send_packet_err );
  PRSTAT(cpriv, n_check_auth_err );
  PRSTAT(cpriv, n_create_challenge_err );
  PRSTAT(cpriv, n_get_challenge_err );
  PRSTAT(cpriv, n_get_response_err );
  PRSTAT(cpriv, n_check_response_err );
  PRSTAT(cpriv, n_check_packet_err );

  fflush(stdout);
}

/*
 * aobj->privateData    : per class data
 * aconn->securityData  : per connection data
 */

/*
 * this is called at connection termination
 * if no connection references the security object
 * than it can be released
 */
RT0D4_CDECL
int test_Close( struct rx_securityClass *aobj )
{
  secobj_priv * sepr = (secobj_priv *)aobj->privateData;
  aobj->refCount--;
  assert( aobj->refCount >= 0 );
  assert( sepr != NULL );

  /* free aobj->privateData */
  /* won't free aobj: it is static */
  if( aobj->refCount == 0 ) delete sepr;
  return 0;
}

/*
 * client: this is created when rx_NewConnection is called
 * server: this is only called when the client actually makes a call
 */
RT0D4_CDECL
int test_NewConnection(
                struct rx_securityClass *aobj,
                struct rx_connection *aconn )
{
  /* allocate aconn->securityData */
  assert( aconn->securityData == 0 );
  assert( aobj->privateData );

  aconn->securityData =
    new conn_priv(
      *((secobj_priv *)aobj->privateData) );

  crypt_buf cb;
  rx_SetSecurityHeaderSize( aconn, (cb.get_header_len() + cb.get_mac_len()) );
  rx_SetSecurityMaxTrailerSize( aconn, 0 );

  aobj->refCount++;
  print_rx_debug( 0, aconn, "test_NewConnection" );
  return 0;
}

RT0D4_CDECL
int test_DestroyConnection(
                struct rx_securityClass *aobj,
                struct rx_connection *aconn )
{
  print_rx_debug( 0, aconn, "test_DestroyConnection" );

  /* free aconn->securityData() */
  if( aconn->securityData ) delete (conn_priv *)aconn->securityData;
  aconn->securityData = 0;

  /* adjust refcount */
  aobj->refCount--;
  assert( aobj->refCount >= 0 );
  return 0;
}

/*
 * only called on the server side to create a challenge
 * when a new client connects (or when it gets unauthenticated ?)
 */
RT0D4_CDECL
int test_CreateChallenge(
                struct rx_securityClass * aobj,
                struct rx_connection * aconn)
{
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  cpriv->set_authenticated(false);
  assert( cpriv->create_challenge() == true );

  cpriv->n_create_challenge_ok(true);
  print_rx_debug( 0, aconn, "test_CreateChallenge" );

  return 0;
}

/*
 * this is the function that actually sends out the challenge data.
 * the challenge data must fit into one iovec size, that is based on the
 * MTU. challenge cannot be split
 */
RT0D4_CDECL
int test_GetChallenge(
                struct rx_securityClass * aobj,
                struct rx_connection * aconn,
                struct rx_packet * apacket)
{
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  size_t challenge_len = 0;
  const unsigned char * challenge = cpriv->get_challenge(challenge_len);

  assert( challenge_len > 0 );
  assert( challenge != NULL );

  rx_packetwrite( apacket,
    0,
    challenge_len,
    (char *)challenge );

  rx_SetDataSize(apacket, challenge_len);

  cpriv->n_get_challenge_ok(true);
  print_rx_debug( apacket, aconn, "test_GetChallenge" );
  printf("    ... Challenge size: %d\n", challenge_len );

  return 0;
}

/*
 * this is called on the server side when a response arrives
 * from the client
 */
RT0D4_CDECL
int test_CheckResponse(
                struct rx_securityClass * aobj,
                struct rx_connection * aconn,
                struct rx_packet * apacket)
{
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  /* push response in */
  assert( cpriv->set_response(                             /* expect to have a contigous */
                (unsigned char *)rx_DataOf(apacket),       /* response packet            */
                (size_t)rx_Contiguous(apacket)) == true );

  cpriv->set_authenticated( true );

  cpriv->n_check_response_ok(true);
  print_rx_debug( apacket, aconn, "test_CheckResponse" );
  return 0;
}

/*
 * this called zillion times on the server when communication
 * takes place with the client
 */
RT0D4_CDECL
int test_CheckAuthentication(
                struct rx_securityClass * aobj,
                struct rx_connection * aconn)
{
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );
  int ret = 0;

  if( cpriv->get_authenticated() )
  {
    cpriv->n_check_auth_ok(true);
    ret = 0;
  }
  else
  {
    cpriv->n_check_auth_err(true);
    ret = 1;
  }
  print_rx_debug( 0, aconn, "test_CheckAuthentication" );
  return ret;
}

/*
 * this encrypt helper only encrypts the first packet
 * if no session key can be found, than it is an error
 */
int encrypt_packet (
                struct rx_securityClass * aobj,
                struct rx_call * acall,
                struct rx_packet * apacket )
{
  struct rx_connection * aconn = (struct rx_connection *)rx_ConnectionOf(acall);
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  if( cpriv->get_session_key() == NULL ) return -1;

  crypt_buf cb;

  unsigned char * data = (unsigned char *)(apacket->wirevec[1].iov_base);
  /* this includes the security header */
  int             sz   = rx_GetDataSize(apacket) + rx_GetSecurityHeaderSize(aconn);
  int             szi  = apacket->wirevec[1].iov_len;
  int             szt  = 0;
  size_t          hl   = cb.get_header_len();
  size_t          ml   = cb.get_mac_len();

  /* init header */
  cb.init_crypt( data, cpriv->get_session_key(), true );
  data += (hl+ml);
  sz  -= (hl+ml);
  szi -= (hl+ml);

  szt = (sz < szi ? sz : szi);

  /* first packet */
  printf("Encrypting: %d bytes\n",szt);
  cb.add_data( data,szt,true );
  sz -= szt;

  /* next packets */
  for( size_t i=2;i<(apacket->niovecs && sz>0);++i )
  {
    szi  = apacket->wirevec[i].iov_len;
    data = (unsigned char *)apacket->wirevec[i].iov_base;
    printf("Encrypting: %d bytes\n",szi);
    cb.add_data( data, szi, true );
    sz -= szi;
  }

  /* add MAC to the second part of the header */
  data = (unsigned char *)(apacket->wirevec[1].iov_base);
  data += hl;
  cb.finalize(data);

  // rx_SetDataSize( apacket, sz );

  return 0;
}

/*
 * this helper is called to ensure that only encrypted or
 * garbage can be sent out from the server
 *
 * There is a security issue:
 *   the first client packet flags the connection initiation
 *   but at that time session key may not be generated yet
 *   so the prepare function has no key to use for encryption
 *
 *   I couldn't find any way so far to delay the packet sending
 *   and still make the server generate and send the challenge
 *   so I chose to garble the packet instead
 */
int garble_packet (
                struct rx_securityClass * aobj,
                struct rx_call * acall,
                struct rx_packet * apacket )
{
  struct rx_connection * aconn = (struct rx_connection *)rx_ConnectionOf(acall);
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  if( cpriv->get_session_key() == NULL )
  {
    char * data = rx_DataOf(apacket);
    int    sz   = rx_Contiguous(apacket)+rx_GetSecurityHeaderSize(aconn);

    memset(data,0xbd,sz);
    /* small amount of garbage is better */
    rx_SetDataSize(apacket, 1 );
  }

  return 0;
}


/* encrypt */
RT0D4_CDECL
int test_PreparePacket (
                struct rx_securityClass * aobj,
                struct rx_call * acall,
                struct rx_packet * apacket )
{
  struct rx_connection * aconn = (struct rx_connection *)rx_ConnectionOf(acall);
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  print_rx_debug( apacket, aconn, "test_PreparePacket /0", acall );

  if( encrypt_packet( aobj, acall, apacket ) != 0 )
  {
    cpriv->n_prepare_packet_err(true);
    garble_packet( aobj, acall, apacket );
    print_rx_debug( apacket, aconn, "test_PreparePacket (KO)", acall );
    return 1450;
  }
  else
  {
    cpriv->n_prepare_packet_ok(true);
    print_rx_debug( apacket, aconn, "test_PreparePacket /1", acall );
    return 0;
  }
}

/* only used to gather packet statistics */
RT0D4_CDECL
int test_SendPacket(
                struct rx_securityClass * aobj,
                struct rx_call * acall,
                struct rx_packet * apacket)
{
  struct rx_connection * aconn = (struct rx_connection *)rx_ConnectionOf(acall);
  conn_priv * cpriv = (conn_priv *)aconn->securityData;

  if( cpriv->get_authenticated() )
  {
    cpriv->n_send_packet_ok(true);
    print_rx_debug( apacket, aconn, "test_SendPacket OK", acall );
  }
  else if( cpriv->n_prepare_packet_ok() == 0 )
  {
    cpriv->n_send_packet_err(true);
    print_rx_debug( apacket, aconn, "test_SendPacket KO", acall );
  }

  return 0;
}

/* decrypt */
RT0D4_CDECL
int test_CheckPacket (
                struct rx_securityClass * aobj,
                struct rx_call * acall,
                struct rx_packet * apacket )
{
  struct rx_connection * aconn = (struct rx_connection *)rx_ConnectionOf(acall);
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  secobj_priv * sepr = (secobj_priv *)aobj->privateData;
  assert( cpriv != NULL );
  assert( sepr != NULL );

  /* this get called even for the challenge/response!
     important: no decryption needed for that! */

  if( cpriv->n_send_packet_ok() == 0 && sepr->is_server() == false &&
      cpriv->n_get_response_ok() == 0 )
  {
    /* no packet sent out so far, so no response to decrypt */
    cpriv->n_check_packet_ok(true);
    print_rx_debug( apacket, aconn, "test_CheckPacket Challenge/Response", acall );
    return 0;
  }

  if( cpriv->get_session_key() == NULL )
  {
    cpriv->n_check_packet_err(true);
    print_rx_debug( apacket, aconn, "test_CheckPacket KO", acall );
    return 0;
  }

  crypt_buf cb;

  unsigned char * data = (unsigned char *)(apacket->wirevec[1].iov_base);
  /* this includes the security header */
  int             sz   = (rx_GetDataSize(apacket));
  int             szi  = apacket->wirevec[1].iov_len;
  int             szt  = 0;
  int             szv  = szi;
  size_t          hl   = cb.get_header_len();
  size_t          ml   = cb.get_mac_len();
  unsigned char   mac[64];
  unsigned int *  pi   = (unsigned int *)data;

  print_rx_debug( apacket, aconn, "test_CheckPacket /0" );

  if( pi[0] == 0xbdbdbdbd )
  {
    if( pi[1] == 0xbdbdbdbd )
    {
      /* sent without encryption */
      print_rx_debug( apacket, aconn, "test_CheckPacket /1 SENT W/O Encryption", acall );
      return 0;
    }
  }

  /* init header */
  cb.init_crypt( data, cpriv->get_session_key(), false );
  data += (hl+ml);
  sz -= (hl+ml);
  szi -= (hl+ml);

  szt = (sz < szi ? sz : szi);

  /* first packet */
  printf("Decrypting: %d bytes\n",szt);
  cb.add_data( data, szt, false );
  sz -= szt;

  /* next packets */
  for( size_t i=2;i<(apacket->niovecs && sz>0);++i )
  {
    szi  = apacket->wirevec[i].iov_len;
    data = (unsigned char *)apacket->wirevec[i].iov_base;
    printf("Decrypting: %d bytes\n",szi);
    cb.add_data( data, szi, false );
    sz -= szi;
    szv += szi;
  }

  /* fix last iov length */
  if( szv > rx_GetDataSize(apacket) )
  {
    szt = szv-rx_GetDataSize(apacket);
    /* cut surplus */
    apacket->wirevec[apacket->niovecs-1].iov_len -= szt;
  }

  /* check mac */
  data = (unsigned char *)(apacket->wirevec[1].iov_base);
  data += hl;

  assert( sizeof(mac) > ml );
  cb.finalize( mac );

  if( memcmp(mac,data,ml) == 0 )
  {
    cpriv->n_check_packet_ok(true);
    print_rx_debug( apacket, aconn, "test_CheckPacket /1", acall );
  }
  else
  {
    cpriv->n_check_packet_err(true);
    print_rx_debug( apacket, aconn, "test_CheckPacket /1 ERROR", acall );

    printf(" MAC : ");
    for( size_t ii=0; ii<ml ;++ii )
      printf("%.2x", mac[ii] );
    printf("\n");

    printf(" DTA : ");
    for( size_t ii=0; ii<ml ;++ii )
      printf("%.2x", data[hl+ii] );
    printf("\n");

    return 1003;
  }
  // rx_SetDataSize( apacket, rx_GetDataSize(apacket) );
  return 0;
}

static struct rx_securityClass::rx_securityOps server_ops = {
    test_Close,
    test_NewConnection,
    test_PreparePacket,
    0, // int (*op_SendPacket) (struct rx_securityClass * aobj, struct rx_call * acall, struct rx_packet * apacket);
    test_CheckAuthentication,
    test_CreateChallenge,
    test_GetChallenge,
    0, // int (*op_GetResponse) (struct rx_securityClass * aobj, struct rx_connection * aconn, struct rx_packet * apacket);
    test_CheckResponse,
    test_CheckPacket,
    test_DestroyConnection,
    0, // int (*op_GetStats) (struct rx_securityClass * aobj, struct rx_connection * aconn, struct rx_securityObjectStats * astats);
    0, // int (*op_SetConfiguration) (struct rx_securityClass * aobj, struct rx_connection * aconn, rx_securityConfigVariables atype,
       //                             void * avalue, void ** acurrentValue);
    0, // int (*op_Spare2) (void);
    0  // int (*op_Spare3) (void);
};

static struct rx_securityClass server_object = {
  &server_ops,
  0,
  0
};

/*
 * client function that reads the challenge, generates the response and
 * sends it
 */
RT0D4_CDECL
int test_GetResponse(
                struct rx_securityClass * aobj,
                struct rx_connection * aconn,
                struct rx_packet * apacket)
{
  conn_priv * cpriv = (conn_priv *)aconn->securityData;
  assert( cpriv != NULL );

  /* push challenge in */
  assert( cpriv->set_challenge(                            /* expect to have a contigous */
                (unsigned char *)rx_DataOf(apacket),       /* response packet            */
                (size_t)rx_Contiguous(apacket)) == true ); /* TODO: fix this             */

  /* pull response out */
  size_t response_len = 0;
  const unsigned char * response = cpriv->get_response(response_len);

  assert( response_len > 0 );
  assert( response != NULL );

  rx_packetwrite( apacket,
    0,
    response_len,
    (char *)response );

  rx_SetDataSize(apacket, response_len);

  if( cpriv->get_session_key() ) cpriv->set_authenticated( true );

  cpriv->n_get_response_ok(true);
  print_rx_debug( apacket, aconn, "test_GetResponse" );

  return 0;
}

static struct rx_securityClass::rx_securityOps client_ops = {
    test_Close,
    test_NewConnection,
    test_PreparePacket,
    test_SendPacket,
    test_CheckAuthentication,
    0, // int (*op_CreateChallenge) (struct rx_securityClass * aobj, struct rx_connection * aconn);
    0, // int (*op_GetChallenge) (struct rx_securityClass * aobj, struct rx_connection * aconn, struct rx_packet * apacket);
    test_GetResponse,
    0, // int (*op_CheckResponse) (struct rx_securityClass * aobj, struct rx_connection * aconn, struct rx_packet * apacket);
    test_CheckPacket,
    test_DestroyConnection,
    0, // int (*op_GetStats) (struct rx_securityClass * aobj, struct rx_connection * aconn, struct rx_securityObjectStats * astats);
    0, // int (*op_SetConfiguration) (struct rx_securityClass * aobj, struct rx_connection * aconn, rx_securityConfigVariables atype,
       //                             void * avalue, void ** acurrentValue);
    0, // int (*op_Spare2) (void);
    0  // int (*op_Spare3) (void);
};

static struct rx_securityClass client_object = {
  &client_ops,
  0,
  0
};


const char * algs_subset[] = {
  "prime192v3", // rel speed  7 : >80b
  "secp224r1",  // rel speed 10 : >80b
  "prime256v1", // rel speed 13 : 128b
  "secp384r1",  // rel speed 33 : 192b
  "secp521r1",  // rel speed 65 : 256b
  NULL
};


secobj_priv * alloc_default_secpr(bool is_server)
{
  secobj_priv * secpr = new secobj_priv(is_server);
  secpr->set_debug(1000);
  assert( secpr->add_algoritms(algs_subset) == true );
  assert( secpr->get_algorithms() != NULL );
  assert( secpr->set_eckey_strength(192,521) == true );
  assert( secpr->get_eckey_strength_min() == 192 );
  assert( secpr->get_eckey_strength_max() == 521 );
  secpr->encrypt_packets(false);
  assert( secpr->encrypt_packets() == false );
  secpr->encrypt_packets(true);
  assert( secpr->encrypt_packets() == true );
  // TODO : cert and credentials
  // TODO : callbacks
  assert( secpr->is_server() == is_server );
  assert( secpr->do_debug() == true );
  return secpr;
}


RT0D4_CDECL
struct rx_securityClass * get_client_security( )
{
  client_object.privateData = alloc_default_secpr(false);
  return &client_object;
}

RT0D4_CDECL
struct rx_securityClass * get_server_security( )
{
  server_object.privateData = alloc_default_secpr(true);
  return &server_object;
}

RT0D4_CDECL
bool validate_client_security_object( struct rx_securityClass * obj )
{
  return true;
}

/* EOF */
